package RC4;

import org.apache.commons.codec.DecoderException;
import org.apache.commons.codec.binary.Hex;

public class RC4 {
    private byte[] S; // 状态数组
    private byte[] T; // 临时数组
    private byte[] Ks; // 密钥流

    public RC4() {
        S = new byte[256];
        T = new byte[256];
    }

    // 密钥编排算法
    private void KSA(byte[] key) {
        // 密钥key的长度在 1 ~ 256 字节之间
        if (key == null || key.length == 0 || key.length > 256) {
            throw new RuntimeException("密钥长度不合理。");
        }

        // 初始化 S 和 T
        int keylen = key.length;
        for (int i = 0; i < 256; i++) {
            S[i] = (byte) i;
            T[i] = key[i % keylen];
        }
        // S 的初始置换
        int j = 0;
        for (int i = 0; i < 256; i++) {
            j = (j + Byte.toUnsignedInt(S[i]) + Byte.toUnsignedInt(T[i])) % 256;
            // 交换 S[i] 和 S[j]
            byte temp = S[i];
            S[i] = S[j];
            S[j] = temp;
        }
    }

    // 密钥流的生成
    private void generateKeys(int ptLength) {
        int i = 0, j = 0;
        Ks = new byte[ptLength];
        int ksIndex = 0;
        while (ksIndex < ptLength) {
            i = (i + 1) % 256;
            j = (j + Byte.toUnsignedInt(S[i])) % 256;

            byte temp = S[i];
            S[i] = S[j];
            S[j] = temp;

            int t = (Byte.toUnsignedInt(S[i]) + Byte.toUnsignedInt(S[j])) % 256;
            Ks[ksIndex++] = S[t];
        }
    }

    public String encrypt(String plaintext, String key) {
        byte[] keyBytes = null;
        byte[] ptBytes = null;
        try {
            keyBytes = Hex.decodeHex(key);
            ptBytes = Hex.decodeHex(plaintext);
        } catch (DecoderException e) {
            throw new RuntimeException(e);
        }
        KSA(keyBytes);
        generateKeys(plaintext.length());

        for (int i = 0; i < ptBytes.length; i++) {
            ptBytes[i] ^= Ks[i];
        }

        return Hex.encodeHexString(ptBytes);
    }

    public String decrypt(String ciphertext, String key) {

        byte[] keyBytes = null;
        byte[] ctBytes = null;
        try {
            keyBytes = Hex.decodeHex(key);
            ctBytes = Hex.decodeHex(ciphertext);
        } catch (DecoderException e) {
            throw new RuntimeException(e);
        }
        KSA(keyBytes);
        generateKeys(ciphertext.length());

        for (int i = 0; i < ctBytes.length; i++) {
            ctBytes[i] ^= Ks[i];
        }

        return Hex.encodeHexString(ctBytes);
    }


}
